from flask import Blueprint,url_for, render_template, redirect, request, jsonify ,g, flash
from decorators import login_required
from exts import db
from models import ProjectsModel, SiteModel, SpeciesModel
import json, math
from decimal import Decimal
from singlesite3pg import calcsinglesite
from multisite3pg import default_parameters, SiteData, parse_age, ThreePGModel


bp = Blueprint("calc", __name__, url_prefix="/calc")

@bp.route("/entersinglesite/<int:projectId>", methods=['GET','POST'])
@login_required
def EnterSingleSite(projectId):
    project = []
    if request.method == 'GET':
        project = ProjectsModel.query.get(projectId)
        dataInput = json.loads(project.data_input)
        dataOutput = []
        if project.data_output:
            dataOutput = json.loads(project.data_output)
        context = {
            "project": project,
            "dataInput": dataInput,
            "dataOutput": dataOutput
        }
        return render_template('singlesite_new.html', **context)
    else:
        pass

@bp.route('/save_singlesite_data', methods=['POST'])
@login_required
def save_singlesite_data():
    try:
        # 获取前端传来的 JSON 数据
        data = request.get_json()
        # print(data)
        project_id = data['projectId']  # 前端传来的项目 ID
        form_data = data['data']  # 前端传来的表单数据
        # 查询数据库中对应的项目
        project = ProjectsModel.query.filter_by(id=project_id).first()
        if not project:
            return jsonify({'success': False, 'message': 'Project not found'}), 404
        # 构造 data_input 的 JSON 结构
        def safe_float(value, default=0):
            try:
                return float(value) if value else default
            except (ValueError, TypeError):
                return default
        def safe_int(value, default=0):
            try:
                return int(value) if value else default
            except (ValueError, TypeError):
                return default
        data_input = {
            "Site": form_data['site'],
            "Species": form_data['species'],
            "Latitude": safe_float(form_data['siteFactors'].get('latitude')),
            "Fertility_rating": safe_float(form_data['siteFactors'].get('fertilityRating')),
            "Soil_class": form_data['siteFactors']['soilClass'],
            "Maximum_ASW": safe_float(form_data['siteFactors'].get('maxASW')),
            "Minimum_ASW": safe_float(form_data['siteFactors'].get('minASW')),
            "Atmospheric_CO2": safe_float(form_data['siteFactors'].get('atmosphericCO2')),
            "Initial_biomass_pools": form_data['initialization']['initialBiomassPools'],
            "Initial_WF": safe_float(form_data['initialization'].get('initialWF')),
            "Initial_WR": safe_float(form_data['initialization'].get('initialWR')),
            "Initial_WS": safe_float(form_data['initialization'].get('initialWS')),
            "Stand_mass": safe_float(form_data['initialization'].get('standMass')),
            "WF_fraction": safe_float(form_data['initialization'].get('wfFraction')),
            "WR_fraction": safe_float(form_data['initialization'].get('wrFraction')),
            "WS_fraction": safe_float(form_data['initialization'].get('wsFraction')),
            "Seedling_mass": safe_float(form_data['initialization'].get('seedlingMass')),
            "Initial_date": form_data['initialization'].get('initialDateOption', ''),
            "Initial_year": safe_int(form_data['initialization'].get('initialYear')),
            "Initial_month": safe_int(form_data['initialization'].get('initialMonth')),
            "Initial_age_y": safe_int(form_data['initialization'].get('initialAgeY')),
            "Initial_age_m": safe_int(form_data['initialization'].get('initialAgeM')),
            "Year_planted": safe_int(form_data['initialization'].get('yearPlanted')),
            "Month_planted": form_data['initialization']['monthPlanted'],
            "Initial_stocking": safe_float(form_data['initialization'].get('initialStocking')),
            "Initial_ASW": safe_float(form_data['initialization'].get('initialASW')),
            "End_age": safe_int(form_data['initialization'].get('endAge')),
            "Interpolate_lookups": form_data['silviculturalEvents'].get('interpolateLookups', 'N'),
            "Fertility": [
                {"Age": safe_float(row.get('age')), "FR": safe_float(row.get('fr'))}
                for row in form_data['silviculturalEvents']['fertility']
            ],
            "Irrigation": [
                {"Age": safe_float(row.get('age')), "Ml_ha_yr": safe_float(row.get('ml_ha_yr', row.get('ml/ha/yr')))}
                for row in form_data['silviculturalEvents']['irrigation']
            ],
            "Thinning": [
                {
                    "Age": safe_float(row.get('age')),
                    "Stocking": safe_float(row.get('stocking')),
                    "F": safe_float(row.get('f')),
                    "R": safe_float(row.get('r')),
                    "S": safe_float(row.get('s'))
                }
                for row in form_data['silviculturalEvents']['thinning']
            ],
            "Defoliation": [
                {"Age": safe_float(row.get('age')), "FR": safe_float(row.get('fr'))}
                for row in form_data['silviculturalEvents']['defoliation']
            ]
        }
        project.data_input = json.dumps(data_input)
        db.session.commit()

        dataOutput = calcsinglesite(project_id)

        project.data_output = json.dumps(dataOutput)
        db.session.commit()

        return jsonify({'success': True, 'message': 'Data saved successfully', 'dataOutput': dataOutput})

    except Exception as e:
        print(f"Error: {str(e)}")
        db.session.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500

@bp.route("/entermultisite/<int:projectId>", methods=['GET', 'POST'])
@login_required
def EnterMultiSite(projectId):
    project = ProjectsModel.query.get(projectId) if projectId else None
    sites = SiteModel.query.filter_by(user_id=g.user.id).all()
    species_list = SpeciesModel.query.filter_by(user_id=g.user.id).all()
    data_input = json.loads(project.data_input) if project and project.data_input else {}

    if request.method == 'POST':
        print("Processing Multi-Site Form Submission")
        form_data = request.form
        print("zhes",form_data)

        # 获取通用设置
        projectname = form_data.get('projectname')
        species = form_data.get('species')
        end_age = int(form_data.get('end_age'))
        output_data = [
            "stemNo", "WF", "WR", "WS", "WL", "standVol", "LAI", "avDBH", "Irrig", "incrWS",
            "mortality", "selfThin", "meantemp", "FR", "CO2", "minASW", "StandAge", "BasArea",
            "Height", "MAI", "MAIx", "ageMAIx", "SLA", "CanCover", "LAIx", "ageLAIx", "TotalW",
            "AvStemMass", "Density", "fracBB", "fAge", "fVPD", "fTemp", "fCalpha", "fCg",
            "fFrost", "fSW", "fNutr", "PhysMod", "GPP", "NPP", "RADint", "alphaC", "Epsilon",
            "CVI", "m", "pR", "pS", "pF", "pFS", "gammaF", "wSmax", "gammaN", "supIrrig",
            "RunOff", "fRainInt", "RainInt", "CanCond", "WUE", "ET", "Transp", "ASW",
            "DayLength", "FrostDays", "SolarRad", "Tav", "VPD", "Rain", "ltStemGR", "Litter",
            "TotalLitter", "EpsilonStem", "pooledSW", "fT"
        ]

        # 确保没有空值（这里已经都是硬编码的非空值，但保留清理逻辑以防将来修改）
        output_data = [var.strip() for var in output_data if var.strip()]

        # 可选字段处理
        month_planted = form_data.get('month_planted')  # 可为空
        # 在表单处理部分（例如 Flask 视图函数）
        output_frequency = form_data.get('output_frequency', 'm')  # 默认值为 "m"
        if output_frequency not in ['m', 'a']:
            output_frequency = 'm'  # 如果输入无效，强制设为 "m"
        interpolate_lookups = form_data.get('interpolate_lookups') == 'True' if form_data.get(
            'interpolate_lookups') else False  # 默认 False

        # 处理站点数据
        sites_data = []
        site_indices = set()
        for key in form_data.keys():
            if key.startswith('sites['):
                index = int(key.split('[')[1].split(']')[0])
                site_indices.add(index)

        species_model = SpeciesModel.query.filter_by(speciesname=species, user_id=g.user.id).first()
        if not species_model:
            flash(f"Species {species} not found.", "error")
            return redirect(url_for('calc.EnterMultiSite', projectId=project.id))

        param_mapping = {
            "SWconst": "SWconst0", "SWpower": "SWpower0", "Y": "y", "aS": "aWs", "nS": "nWs",
            "alpha": "alphaCx", "gammaNx": "gammaN1", "gammaFx": "gammaF1", "rhoMin": "rho0", "rhoMax": "rho1"
        }
        species_params = default_parameters.copy()
        exclude_fields = {'id', 'speciesname', 'user_id', 'create_time'}
        for key, value in species_model.__dict__.items():
            mapped_key = param_mapping.get(key, key)
            if mapped_key in species_params and key not in exclude_fields:
                species_params[mapped_key] = float(value) if isinstance(value, Decimal) else value

        for i in sorted(site_indices):
            site = SiteData()
            site.Sites = form_data.get(f'sites[{i}][sitename]')  # 必填字段

            # 可选字段，带默认值或跳过
            site.Latitude = float(form_data.get(f'sites[{i}][latitude]', 0.0)) if form_data.get(
                f'sites[{i}][latitude]') else 0.0
            site.FR = float(form_data.get(f'sites[{i}][fr]', 1.0)) if form_data.get(f'sites[{i}][fr]') else 1.0
            site.Soil_type = form_data.get(f'sites[{i}][soil_type]', 'CL')
            site.Max_ASW = float(form_data.get(f'sites[{i}][max_asw]', 200.0)) if form_data.get(
                f'sites[{i}][max_asw]') else 200.0
            site.Min_ASW = float(form_data.get(f'sites[{i}][min_asw]', 0.0)) if form_data.get(
                f'sites[{i}][min_asw]') else 0.0
            site.Initial_stocking = float(form_data.get(f'sites[{i}][initial_stocking]', 1000.0)) if form_data.get(
                f'sites[{i}][initial_stocking]') else 1000.0
            site.gammaNx = float(form_data.get(f'sites[{i}][gammaNx]', 1.0)) if form_data.get(
                f'sites[{i}][gammaNx]') else 1.0
            site.alphaCx = float(form_data.get(f'sites[{i}][alphaCx]', 0.06)) if form_data.get(
                f'sites[{i}][alphaCx]') else 0.06  # 动态输入，默认 0.06
            planted_year = int(form_data.get(f'sites[{i}][planted_year]', 1990)) if form_data.get(
                f'sites[{i}][planted_year]') else 1990  # 动态输入，默认 1990
            month_planted_value = month_planted if month_planted else "June"

            sites_data.append({
                'site': site,
                'planted_year': planted_year,
                'month_planted': month_planted_value,
                'end_age': end_age
            })

        # 林务事件处理（保持不变）
        silvicultural_events = {}
        for event_type in ['thinning', 'defoliation', 'fertility', 'irrigation']:
            event_indices = set()
            for key in form_data.keys():
                if key.startswith(f'{event_type}['):
                    index = int(key.split('[')[1].split(']')[0])
                    event_indices.add(index)

            if event_indices:
                ages = []
                event_data = {}
                for i in sorted(event_indices):
                    age_str = form_data.get(f'{event_type}[{i}][age]')
                    if not age_str:
                        continue
                    try:
                        age_years = parse_age(age_str)
                        ages.append(age_years)
                    except ValueError:
                        continue

                    if event_type == 'thinning':
                        event_data.setdefault('residual_stems', []).append(
                            float(form_data.get(f'{event_type}[{i}][stems]', 0)))
                        event_data.setdefault('thinWF', []).append(
                            float(form_data.get(f'{event_type}[{i}][wf_frac]', 0)))
                        event_data.setdefault('thinWR', []).append(
                            float(form_data.get(f'{event_type}[{i}][wr_frac]', 0)))
                        event_data.setdefault('thinWS', []).append(
                            float(form_data.get(f'{event_type}[{i}][ws_frac]', 0)))
                    elif event_type == 'defoliation':
                        event_data.setdefault('defol_fraction', []).append(
                            float(form_data.get(f'{event_type}[{i}][foliage_remaining]', 0)))
                    elif event_type == 'fertility':
                        event_data.setdefault('FR', []).append(float(form_data.get(f'{event_type}[{i}][fr]', 0)))
                    elif event_type == 'irrigation':
                        event_data.setdefault('irrig', []).append(
                            float(form_data.get(f'{event_type}[{i}][irrigation]', 0)))

                if ages:
                    event_data['ages'] = ages
                    silvicultural_events[event_type.capitalize()] = event_data

        # 保存项目数据
        if project:
            project.projectname = projectname
            project.projecttype = 'MultiSite'
        else:
            project = ProjectsModel(projectname=projectname, projecttype='MultiSite', user_id=g.user.id)
            db.session.add(project)

        data_input = {
            'Species': species,
            'End_age': end_age,
            'Output_data': ','.join(output_data),
            'Sites': [
                {
                    'sitename': s['site'].Sites,
                    'Latitude': s['site'].Latitude,
                    'FR': s['site'].FR,
                    'Soil_class': s['site'].Soil_type,
                    'maxASW': s['site'].Max_ASW ,
                    'minASW': s['site'].Min_ASW,
                    'Initial_stocking': s['site'].Initial_stocking ,
                    'gammaNx': s['site'].gammaNx,
                    'alphaCx': s['site'].alphaCx,
                    'planted_year': s['planted_year'],
                    'date_planted': f"{s['planted_year']}&{s['month_planted'][:3]}",
                    'initial_date': f"{s['planted_year']}&{s['month_planted'][:3]}",
                    'end_age': s['end_age']
                } for s in sites_data
            ]
        }

        # 添加可选字段到 data_input
        if month_planted:
            data_input['Month_planted'] = month_planted
        if form_data.get('output_frequency'):  # 只有勾选并提交时才保存
            data_input['Output_frequency'] = output_frequency
        if form_data.get('interpolate_lookups'):
            data_input['Interpolate_lookups'] = 'True' if form_data.get('interpolate_lookups') == 'Yes' else 'False'

        # 保存林务事件
        for event_type, events in silvicultural_events.items():
            data_input[event_type] = [
                dict(zip(['Age'] + list(events.keys()), [age] + [events[k][i] for k in events.keys() if k != 'ages']))
                for i, age in enumerate(events['ages'])
            ]

        project.data_input = json.dumps(data_input)

        # 执行计算（保持不变）
        results = {}
        for site_data in sites_data:
            site = site_data['site']
            db_site = SiteModel.query.filter_by(sitename=site.Sites, user_id=g.user.id).first()
            Tas_metdata = {
                "TMAX1": [
                    float(db_site.tmax1), float(db_site.tmax2), float(db_site.tmax3),
                    float(db_site.tmax4), float(db_site.tmax5), float(db_site.tmax6),
                    float(db_site.tmax7), float(db_site.tmax8), float(db_site.tmax9),
                    float(db_site.tmax10), float(db_site.tmax11), float(db_site.tmax12)
                ],
                "TMIN1": [
                    float(db_site.tmin1), float(db_site.tmin2), float(db_site.tmin3),
                    float(db_site.tmin4), float(db_site.tmin5), float(db_site.tmin6),
                    float(db_site.tmin7), float(db_site.tmin8), float(db_site.tmin9),
                    float(db_site.tmin10), float(db_site.tmin11), float(db_site.tmin12)
                ],
                "PRECIP1": [
                    float(db_site.precip1), float(db_site.precip2), float(db_site.precip3),
                    float(db_site.precip4), float(db_site.precip5), float(db_site.precip6),
                    float(db_site.precip7), float(db_site.precip8), float(db_site.precip9),
                    float(db_site.precip10), float(db_site.precip11), float(db_site.precip12)
                ],
                "EVAP1": [
                    float(db_site.evap1), float(db_site.evap2), float(db_site.evap3),
                    float(db_site.evap4), float(db_site.evap5), float(db_site.evap6),
                    float(db_site.evap7), float(db_site.evap8), float(db_site.evap9),
                    float(db_site.evap10), float(db_site.evap11), float(db_site.evap12)
                ],
                "RADN1": [
                    float(db_site.ra1), float(db_site.ra2), float(db_site.ra3),
                    float(db_site.ra4), float(db_site.ra5), float(db_site.ra6),
                    float(db_site.ra7), float(db_site.ra8), float(db_site.ra9),
                    float(db_site.ra10), float(db_site.ra11), float(db_site.ra12)
                ],
                "RAIND1": [
                    float(db_site.raind1), float(db_site.raind2), float(db_site.raind3),
                    float(db_site.raind4), float(db_site.raind5), float(db_site.raind6),
                    float(db_site.raind7), float(db_site.raind8), float(db_site.raind9),
                    float(db_site.raind10), float(db_site.raind11), float(db_site.raind12)
                ],
                "FROST1": [
                    float(db_site.frost1), float(db_site.frost2), float(db_site.frost3),
                    float(db_site.frost4), float(db_site.frost5), float(db_site.frost6),
                    float(db_site.frost7), float(db_site.frost8), float(db_site.frost9),
                    float(db_site.frost10), float(db_site.frost11), float(db_site.frost12)
                ]
            }

            model = ThreePGModel(
                site_data=site,
                species=species,
                climate_data=Tas_metdata,
                seedling_mass=1,
                month_planted=site_data['month_planted'],
                planted_year=site_data['planted_year'],
                end_age=site_data['end_age'],
                lookup=interpolate_lookups,
                params=species_params,
                output_frequency=output_frequency,
                silvicultural_events=silvicultural_events
            )
            site_results = model.run3PG(output_vars=output_data)
            results[site.Sites] = site_results
        project.data_output = json.dumps(results)
        db.session.commit()

        flash('Multi-Site run completed successfully!', 'success')
        return render_template('multisite_results_new.html', project=project, data_output=results, output_vars=output_data)

    return render_template(
        'multisite_new.html',
        project=project,
        sites=sites,
        species_list=species_list,
        data_input=data_input
    )

@bp.route("/entersensitivity/<int:projectId>", methods=['GET','POST'])
@login_required
def EnterSensitivity(projectId):
    project = []
    if request.method == 'GET':
        project = ProjectsModel.query.get(projectId)
        dataInput = json.loads(project.data_input)
        if "Benchmark" in dataInput:
            Benchmark = ProjectsModel.query.get(dataInput["Benchmark"])
        else:
            Benchmark = []
        singlesiteProjects = ProjectsModel.query.filter_by(projecttype='SingleSite', user_id=g.user.id).all()
        return render_template('sensitivity.html', project=project, singlesiteProjects=singlesiteProjects,Benchmark=Benchmark)
    else:
        pass